package com.atguigu.app.func;

import com.alibaba.fastjson.JSON;
import com.alibaba.fastjson.JSONObject;
import com.atguigu.common.Constant;
import com.atguigu.util.HBaseUtil;
import com.atguigu.util.JedisUtil;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.api.async.RedisAsyncCommands;
import lombok.SneakyThrows;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.streaming.api.functions.async.ResultFuture;
import org.apache.flink.streaming.api.functions.async.RichAsyncFunction;
import org.apache.hadoop.hbase.client.AsyncConnection;

import java.util.Collections;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;
import java.util.function.Function;
import java.util.function.Supplier;

public abstract class DimInfoAsyncFunction<T> extends RichAsyncFunction<T, T> {

    private AsyncConnection asyncConnection;
    private StatefulRedisConnection<String, String> redisConnection;
    private String tableName;

    public DimInfoAsyncFunction(String tableName) {
        this.tableName = tableName;
    }

    @Override
    public void open(Configuration parameters) throws Exception {
        asyncConnection = HBaseUtil.getAsyncConnection();
        redisConnection = JedisUtil.getAsyncRedisConnection();
    }

    public abstract String getPk(T input);

    public abstract void join(T input, JSONObject dimInfo);

    @Override
    public void asyncInvoke(T input, ResultFuture<T> resultFuture) throws Exception {
        CompletableFuture.supplyAsync(new Supplier<JSONObject>() {
            @SneakyThrows
            @Override
            public JSONObject get() {
                RedisAsyncCommands<String, String> async = redisConnection.async();
                String pk = getPk(input);
                String redisKey = "DIM:" + tableName + ":" + pk;
                String dimInfoStr = async.get(redisKey).get();
                if (dimInfoStr != null) {
                    return JSON.parseObject(dimInfoStr);
                } else {
                    return null;
                }
            }
        }).thenApplyAsync(new Function<JSONObject, JSONObject>() {
            @SneakyThrows
            @Override
            public JSONObject apply(JSONObject jsonObject) {
                if (jsonObject != null) { //在Redis已经查询到了数据
                    return jsonObject;
                } else { //在Redis没有查询到数据,则访问HBase,并将数据写到Redis
                    String pk = getPk(input);
                    JSONObject dimInfoJson = HBaseUtil.getData(asyncConnection, Constant.HBASE_NAME_SPACE, tableName, pk);
                    JedisUtil.setDataByAsync(redisConnection, tableName, pk, dimInfoJson.toJSONString());
                    return dimInfoJson;
                }
            }
        }).thenAccept(new Consumer<JSONObject>() {
            @Override
            public void accept(JSONObject jsonObject) {
                //将维表信息补充到数据上
                join(input, jsonObject);
                resultFuture.complete(Collections.singletonList(input));
            }
        });
    }

    @Override
    public void timeout(T input, ResultFuture<T> resultFuture) throws Exception {
        System.out.println("TimeOut>>>>>>>>>" + input);
    }
}
